{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d870f870-e649-46ec-9631-8ae070b01f31",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import os\n",
    "from tempfile import TemporaryDirectory\n",
    "from typing import Tuple\n",
    "\n",
    "import torch\n",
    "from torch import nn, Tensor\n",
    "from torch.nn import TransformerEncoder, TransformerEncoderLayer\n",
    "from torch.utils.data import dataset\n",
    "\n",
    "class TransformerModel(nn.Module):\n",
    "\n",
    "    def __init__(self, ntoken: int, d_model: int, nhead: int, d_hid: int,\n",
    "                 nlayers: int, dropout: float = 0.5):\n",
    "        super().__init__()\n",
    "        self.model_type = 'Transformer'\n",
    "        self.pos_encoder = PositionalEncoding(d_model, dropout)\n",
    "        encoder_layers = TransformerEncoderLayer(d_model, nhead, d_hid, dropout)\n",
    "        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)\n",
    "        self.embedding = nn.Embedding(ntoken, d_model)\n",
    "        self.d_model = d_model\n",
    "        self.linear = nn.Linear(d_model, ntoken)\n",
    "\n",
    "        self.init_weights()\n",
    "\n",
    "    def init_weights(self) -> None:\n",
    "        initrange = 0.1\n",
    "        self.embedding.weight.data.uniform_(-initrange, initrange)\n",
    "        self.linear.bias.data.zero_()\n",
    "        self.linear.weight.data.uniform_(-initrange, initrange)\n",
    "\n",
    "    def forward(self, src: Tensor, src_mask: Tensor = None) -> Tensor:\n",
    "        \"\"\"\n",
    "        Arguments:\n",
    "            src: Tensor, shape ``[seq_len, batch_size]``\n",
    "            src_mask: Tensor, shape ``[seq_len, seq_len]``\n",
    "\n",
    "        Returns:\n",
    "            output Tensor of shape ``[seq_len, batch_size, ntoken]``\n",
    "        \"\"\"\n",
    "        src = self.embedding(src) * math.sqrt(self.d_model)\n",
    "        src = self.pos_encoder(src)\n",
    "        output = self.transformer_encoder(src, src_mask)\n",
    "        output = self.linear(output)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d0b62743-5972-40a4-8203-c82f6a999bbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):\n",
    "        super().__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "        position = torch.arange(max_len).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n",
    "        pe = torch.zeros(max_len, 1, d_model)\n",
    "        pe[:, 0, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 0, 1::2] = torch.cos(position * div_term)\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Arguments:\n",
    "            x: Tensor, shape ``[seq_len, batch_size, embedding_dim]``\n",
    "        \"\"\"\n",
    "        x = x + self.pe[:x.size(0)]\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a30df235-10ba-455b-b5cf-412105ae6ceb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2049990])\n",
      "torch.Size([102499, 20])\n",
      "tensor([[    9,    59,   564,  ..., 11652,  2435,     1],\n",
      "        [ 3849,    12,   300,  ...,    47,    30,  1990],\n",
      "        [ 3869,   315,    19,  ...,    97,  7720,     4],\n",
      "        ...,\n",
      "        [  587,  4011,    59,  ...,     1,  1439, 12313],\n",
      "        [ 4987,    29,     4,  ...,  3165, 17106,  2060],\n",
      "        [    6,     8,     1,  ...,    62,    18,     2]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "from torchtext.datasets import WikiText2\n",
    "from torchtext.data.utils import get_tokenizer\n",
    "from torchtext.vocab import build_vocab_from_iterator\n",
    "\n",
    "train_iter = WikiText2(split='train')\n",
    "tokenizer = get_tokenizer('basic_english')\n",
    "vocab = build_vocab_from_iterator(map(tokenizer, train_iter), specials=['<unk>'])\n",
    "vocab.set_default_index(vocab['<unk>'])\n",
    "\n",
    "def data_process(raw_text_iter: dataset.IterableDataset) -> Tensor:\n",
    "    \"\"\"Converts raw text into a flat Tensor.\"\"\"\n",
    "    data = [torch.tensor(vocab(tokenizer(item)), dtype=torch.long) for item in raw_text_iter]\n",
    "    return torch.cat(tuple(filter(lambda t: t.numel() > 0, data)))\n",
    "\n",
    "# ``train_iter`` was \"consumed\" by the process of building the vocab,\n",
    "# so we have to create it again\n",
    "train_iter, val_iter, test_iter = WikiText2()\n",
    "train_data = data_process(train_iter)\n",
    "val_data = data_process(val_iter)\n",
    "test_data = data_process(test_iter)\n",
    "\n",
    "print(train_data.shape)\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "def batchify(data: Tensor, bsz: int) -> Tensor:\n",
    "    \"\"\"Divides the data into ``bsz`` separate sequences, removing extra elements\n",
    "    that wouldn't cleanly fit.\n",
    "\n",
    "    Arguments:\n",
    "        data: Tensor, shape ``[N]``\n",
    "        bsz: int, batch size\n",
    "\n",
    "    Returns:\n",
    "        Tensor of shape ``[N // bsz, bsz]``\n",
    "    \"\"\"\n",
    "    seq_len = data.size(0) // bsz\n",
    "    data = data[:seq_len * bsz]\n",
    "    data = data.view(bsz, seq_len).t().contiguous()\n",
    "    return data.to(device)\n",
    "\n",
    "batch_size = 20\n",
    "eval_batch_size = 10\n",
    "train_data = batchify(train_data, batch_size)  # shape ``[seq_len, batch_size]``\n",
    "val_data = batchify(val_data, eval_batch_size)\n",
    "test_data = batchify(test_data, eval_batch_size)\n",
    "print(train_data.shape)\n",
    "print(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d7b4f49f-acf5-4c18-b983-51e1049eada4",
   "metadata": {},
   "outputs": [],
   "source": [
    "bptt = 35\n",
    "def get_batch(source: Tensor, i: int) -> Tuple[Tensor, Tensor]:\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        source: Tensor, shape ``[full_seq_len, batch_size]``\n",
    "        i: int\n",
    "\n",
    "    Returns:\n",
    "        tuple (data, target), where data has shape ``[seq_len, batch_size]`` and\n",
    "        target has shape ``[seq_len * batch_size]``\n",
    "    \"\"\"\n",
    "    seq_len = min(bptt, len(source) - 1 - i)\n",
    "    data = source[i:i+seq_len]\n",
    "    target = source[i+1:i+1+seq_len].reshape(-1)\n",
    "    return data, target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0162f2a6-fb9e-4183-a894-8285eabc11a6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TransformerModel(\n",
       "  (pos_encoder): PositionalEncoding(\n",
       "    (dropout): Dropout(p=0.2, inplace=False)\n",
       "  )\n",
       "  (transformer_encoder): TransformerEncoder(\n",
       "    (layers): ModuleList(\n",
       "      (0): TransformerEncoderLayer(\n",
       "        (self_attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=200, out_features=200, bias=True)\n",
       "        )\n",
       "        (linear1): Linear(in_features=200, out_features=200, bias=True)\n",
       "        (dropout): Dropout(p=0.2, inplace=False)\n",
       "        (linear2): Linear(in_features=200, out_features=200, bias=True)\n",
       "        (norm1): LayerNorm((200,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((200,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout1): Dropout(p=0.2, inplace=False)\n",
       "        (dropout2): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "      (1): TransformerEncoderLayer(\n",
       "        (self_attn): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=200, out_features=200, bias=True)\n",
       "        )\n",
       "        (linear1): Linear(in_features=200, out_features=200, bias=True)\n",
       "        (dropout): Dropout(p=0.2, inplace=False)\n",
       "        (linear2): Linear(in_features=200, out_features=200, bias=True)\n",
       "        (norm1): LayerNorm((200,), eps=1e-05, elementwise_affine=True)\n",
       "        (norm2): LayerNorm((200,), eps=1e-05, elementwise_affine=True)\n",
       "        (dropout1): Dropout(p=0.2, inplace=False)\n",
       "        (dropout2): Dropout(p=0.2, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (embedding): Embedding(28782, 200)\n",
       "  (linear): Linear(in_features=200, out_features=28782, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ntokens = len(vocab)  # size of vocabulary\n",
    "emsize = 200  # embedding dimension\n",
    "d_hid = 200  # dimension of the feedforward network model in ``nn.TransformerEncoder``\n",
    "nlayers = 2  # number of ``nn.TransformerEncoderLayer`` in ``nn.TransformerEncoder``\n",
    "nhead = 2  # number of heads in ``nn.MultiheadAttention``\n",
    "dropout = 0.2  # dropout probability\n",
    "model = TransformerModel(ntokens, emsize, nhead, d_hid, nlayers, dropout).to(device)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c791ffa0-41fc-4f6d-9cc7-e7c456921749",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "lr = 5.0  # learning rate\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=lr)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)\n",
    "\n",
    "def train(model: nn.Module) -> None:\n",
    "    model.train()  # turn on train mode\n",
    "    total_loss = 0.\n",
    "    log_interval = 200\n",
    "    start_time = time.time()\n",
    "\n",
    "    num_batches = len(train_data) // bptt\n",
    "    for batch, i in enumerate(range(0, train_data.size(0) - 1, bptt)):\n",
    "        data, targets = get_batch(train_data, i)\n",
    "        output = model(data)\n",
    "        output_flat = output.view(-1, ntokens)\n",
    "        loss = criterion(output_flat, targets)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)\n",
    "        optimizer.step()\n",
    "\n",
    "        total_loss += loss.item()\n",
    "        if batch % log_interval == 0 and batch > 0:\n",
    "            lr = scheduler.get_last_lr()[0]\n",
    "            ms_per_batch = (time.time() - start_time) * 1000 / log_interval\n",
    "            cur_loss = total_loss / log_interval\n",
    "            ppl = math.exp(cur_loss)\n",
    "            print(f'| epoch {epoch:3d} | {batch:5d}/{num_batches:5d} batches | '\n",
    "                  f'lr {lr:02.2f} | ms/batch {ms_per_batch:5.2f} | '\n",
    "                  f'loss {cur_loss:5.2f} | ppl {ppl:8.2f}')\n",
    "            total_loss = 0\n",
    "            start_time = time.time()\n",
    "\n",
    "def evaluate(model: nn.Module, eval_data: Tensor) -> float:\n",
    "    model.eval()  # turn on evaluation mode\n",
    "    total_loss = 0.\n",
    "    with torch.no_grad():\n",
    "        for i in range(0, eval_data.size(0) - 1, bptt):\n",
    "            data, targets = get_batch(eval_data, i)\n",
    "            seq_len = data.size(0)\n",
    "            output = model(data)\n",
    "            output_flat = output.view(-1, ntokens)\n",
    "            total_loss += seq_len * criterion(output_flat, targets).item()\n",
    "    return total_loss / (len(eval_data) - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0e4b1bf0-ca2e-4c37-ac5a-12b279108572",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tempdir: /tmp/tmpnth1zkzl\n",
      "| epoch   1 |   200/ 2928 batches | lr 5.00 | ms/batch 10.20 | loss  4.36 | ppl    78.05\n",
      "| epoch   1 |   400/ 2928 batches | lr 5.00 | ms/batch  9.96 | loss  4.00 | ppl    54.58\n",
      "| epoch   1 |   600/ 2928 batches | lr 5.00 | ms/batch 11.61 | loss  3.66 | ppl    38.98\n",
      "| epoch   1 |   800/ 2928 batches | lr 5.00 | ms/batch 11.32 | loss  3.52 | ppl    33.73\n",
      "| epoch   1 |  1000/ 2928 batches | lr 5.00 | ms/batch 11.27 | loss  3.35 | ppl    28.55\n",
      "| epoch   1 |  1200/ 2928 batches | lr 5.00 | ms/batch 13.18 | loss  3.40 | ppl    30.02\n",
      "| epoch   1 |  1400/ 2928 batches | lr 5.00 | ms/batch 11.26 | loss  3.46 | ppl    31.69\n",
      "| epoch   1 |  1600/ 2928 batches | lr 5.00 | ms/batch  9.60 | loss  3.41 | ppl    30.31\n",
      "| epoch   1 |  1800/ 2928 batches | lr 5.00 | ms/batch  9.24 | loss  3.29 | ppl    26.80\n",
      "| epoch   1 |  2000/ 2928 batches | lr 5.00 | ms/batch  9.65 | loss  3.23 | ppl    25.29\n",
      "| epoch   1 |  2200/ 2928 batches | lr 5.00 | ms/batch 12.04 | loss  3.14 | ppl    23.07\n",
      "| epoch   1 |  2400/ 2928 batches | lr 5.00 | ms/batch 10.88 | loss  3.16 | ppl    23.47\n",
      "| epoch   1 |  2600/ 2928 batches | lr 5.00 | ms/batch  9.61 | loss  3.10 | ppl    22.30\n",
      "| epoch   1 |  2800/ 2928 batches | lr 5.00 | ms/batch  8.67 | loss  3.03 | ppl    20.80\n",
      "-----------------------------------------------------------------------------------------\n",
      "| end of epoch   1 | time: 31.97s | valid loss  2.17 | valid ppl     8.74\n",
      "-----------------------------------------------------------------------------------------\n",
      "| epoch   2 |   200/ 2928 batches | lr 4.75 | ms/batch  8.21 | loss  2.73 | ppl    15.35\n",
      "| epoch   2 |   400/ 2928 batches | lr 4.75 | ms/batch 12.87 | loss  2.72 | ppl    15.17\n",
      "| epoch   2 |   600/ 2928 batches | lr 4.75 | ms/batch  8.24 | loss  2.57 | ppl    13.11\n",
      "| epoch   2 |   800/ 2928 batches | lr 4.75 | ms/batch  8.16 | loss  2.55 | ppl    12.80\n",
      "| epoch   2 |  1000/ 2928 batches | lr 4.75 | ms/batch  8.54 | loss  2.50 | ppl    12.16\n",
      "| epoch   2 |  1200/ 2928 batches | lr 4.75 | ms/batch  8.55 | loss  2.56 | ppl    12.94\n",
      "| epoch   2 |  1400/ 2928 batches | lr 4.75 | ms/batch  7.50 | loss  2.61 | ppl    13.58\n",
      "| epoch   2 |  1600/ 2928 batches | lr 4.75 | ms/batch 13.12 | loss  2.59 | ppl    13.29\n",
      "| epoch   2 |  1800/ 2928 batches | lr 4.75 | ms/batch 12.96 | loss  2.52 | ppl    12.48\n",
      "| epoch   2 |  2000/ 2928 batches | lr 4.75 | ms/batch  8.60 | loss  2.60 | ppl    13.47\n",
      "| epoch   2 |  2200/ 2928 batches | lr 4.75 | ms/batch 11.87 | loss  2.46 | ppl    11.67\n",
      "| epoch   2 |  2400/ 2928 batches | lr 4.75 | ms/batch 10.10 | loss  2.56 | ppl    12.88\n",
      "| epoch   2 |  2600/ 2928 batches | lr 4.75 | ms/batch 12.89 | loss  2.63 | ppl    13.87\n",
      "| epoch   2 |  2800/ 2928 batches | lr 4.75 | ms/batch 13.42 | loss  2.55 | ppl    12.77\n",
      "-----------------------------------------------------------------------------------------\n",
      "| end of epoch   2 | time: 31.17s | valid loss  1.38 | valid ppl     3.99\n",
      "-----------------------------------------------------------------------------------------\n",
      "| epoch   3 |   200/ 2928 batches | lr 4.51 | ms/batch 10.00 | loss  2.27 | ppl     9.64\n",
      "| epoch   3 |   400/ 2928 batches | lr 4.51 | ms/batch 11.53 | loss  2.29 | ppl     9.86\n",
      "| epoch   3 |   600/ 2928 batches | lr 4.51 | ms/batch 10.95 | loss  2.23 | ppl     9.27\n",
      "| epoch   3 |   800/ 2928 batches | lr 4.51 | ms/batch 11.43 | loss  2.28 | ppl     9.75\n",
      "| epoch   3 |  1000/ 2928 batches | lr 4.51 | ms/batch 12.27 | loss  2.24 | ppl     9.35\n",
      "| epoch   3 |  1200/ 2928 batches | lr 4.51 | ms/batch 11.27 | loss  2.32 | ppl    10.16\n",
      "| epoch   3 |  1400/ 2928 batches | lr 4.51 | ms/batch 11.91 | loss  2.28 | ppl     9.82\n",
      "| epoch   3 |  1600/ 2928 batches | lr 4.51 | ms/batch  8.35 | loss  2.36 | ppl    10.57\n",
      "| epoch   3 |  1800/ 2928 batches | lr 4.51 | ms/batch  7.15 | loss  2.28 | ppl     9.80\n",
      "| epoch   3 |  2000/ 2928 batches | lr 4.51 | ms/batch  9.61 | loss  2.34 | ppl    10.35\n",
      "| epoch   3 |  2200/ 2928 batches | lr 4.51 | ms/batch 13.29 | loss  2.26 | ppl     9.59\n",
      "| epoch   3 |  2400/ 2928 batches | lr 4.51 | ms/batch 13.74 | loss  2.28 | ppl     9.81\n",
      "| epoch   3 |  2600/ 2928 batches | lr 4.51 | ms/batch 12.92 | loss  2.40 | ppl    11.02\n",
      "| epoch   3 |  2800/ 2928 batches | lr 4.51 | ms/batch 12.56 | loss  2.34 | ppl    10.37\n",
      "-----------------------------------------------------------------------------------------\n",
      "| end of epoch   3 | time: 33.37s | valid loss  1.48 | valid ppl     4.41\n",
      "-----------------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "best_val_loss = float('inf')\n",
    "epochs = 3\n",
    "\n",
    "with TemporaryDirectory() as tempdir:\n",
    "    print(\"tempdir:\", tempdir)\n",
    "    best_model_params_path = os.path.join(tempdir, \"best_model_params.pt\")\n",
    "\n",
    "    for epoch in range(1, epochs + 1):\n",
    "        epoch_start_time = time.time()\n",
    "        train(model)\n",
    "        val_loss = evaluate(model, val_data)\n",
    "        val_ppl = math.exp(val_loss)\n",
    "        elapsed = time.time() - epoch_start_time\n",
    "        print('-' * 89)\n",
    "        print(f'| end of epoch {epoch:3d} | time: {elapsed:5.2f}s | '\n",
    "            f'valid loss {val_loss:5.2f} | valid ppl {val_ppl:8.2f}')\n",
    "        print('-' * 89)\n",
    "\n",
    "        if val_loss < best_val_loss:\n",
    "            best_val_loss = val_loss\n",
    "            # 保存模型\n",
    "            torch.save(model.state_dict(), best_model_params_path)\n",
    "\n",
    "        scheduler.step()\n",
    "    # 加载最好的模型    \n",
    "    model.load_state_dict(torch.load(best_model_params_path)) # load best model states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a8206bdb-91d6-4b7b-8bb4-f3a8eecfe26e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========================================================================================\n",
      "| End of training | test loss  1.37 | test ppl     3.93\n",
      "=========================================================================================\n"
     ]
    }
   ],
   "source": [
    "test_loss = evaluate(model, test_data)\n",
    "test_ppl = math.exp(test_loss)\n",
    "print('=' * 89)\n",
    "print(f'| End of training | test loss {test_loss:5.2f} | '\n",
    "      f'test ppl {test_ppl:8.2f}')\n",
    "print('=' * 89)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf8296ca-705b-41aa-a3ce-e69f8aa8e710",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
